from pymilvus import (
    connections, Collection, CollectionSchema, FieldSchema,
    DataType, utility
)
from langchain_milvus import Milvus as MilvusVectorStore
from langchain_ollama import OllamaEmbeddings
from typing import List, Dict, Optional, Any
from app.system_config import settings
from langchain_core.documents import Document
from app.exception.rag_exception import RAGException
import logging
import time  # 新增：确保导入time模块


logger = logging.getLogger("Contract-Milvus")


class MilvusStoreWrapper:
    def __init__(self, is_removed: bool = False,collection_name:str = settings.MILVUS_COLLECTION,max_retries: int = 1, retry_delay: int = 2):
        self.collection_name = collection_name
        self.max_retries = max_retries
        self.retry_delay = retry_delay
        self.is_removed = is_removed
        self.vector_field_name = "embedding"  # 向量字段名常量，便于维护
        self._connect()
        self.embeddings = OllamaEmbeddings(
            model=settings.OLLAMA_EMBEDDING_MODEL,
            base_url=settings.OLLAMA_URL
        )
        self.collection = self._init_collection()
        # 确保集合加载完成
        self.collection.load()
        # 验证字段存在性
        self._verify_fields()
        self.vector_store = self._init_langchain_milvus()

    def _connect(self):
        """连接Milvus，带重试机制"""
        for attempt in range(self.max_retries):
            try:
                connections.connect(
                    alias="default",
                    host=settings.MILVUS_HOST,
                    port=settings.MILVUS_PORT,
                    timeout=30  # 设置超时时间
                )
                logger.info(f"成功连接Milvus：{settings.MILVUS_HOST}:{settings.MILVUS_PORT}")
                return
            except Exception as e:
                logger.warning(f"Milvus连接尝试 {attempt + 1}/{self.max_retries} 失败：{str(e)}")
                if attempt < self.max_retries - 1:
                    time.sleep(self.retry_delay)
        raise RAGException(503, "无法连接到Milvus服务，请检查服务是否运行")

    def _verify_fields(self):
        """验证必要字段是否存在"""
        schema = self.collection.schema
        required_fields = [self.vector_field_name, "Comments", "Trapped", "content", "id"]
        missing_fields = [f for f in required_fields if not any(field.name == f for field in schema.fields)]

        if missing_fields:
            logger.error(f"集合缺少必要字段：{missing_fields}")
            raise RAGException(500, f"集合结构错误，缺少字段：{missing_fields}")
        logger.info("所有必要字段验证通过")

    def _init_collection(self) -> Collection:
        """初始化集合，确保可空字段正确配置"""
        collection_name = self.collection_name
        # print(f"self.is_removed:{self.is_removed}")
        if self.is_removed:
            logger.warning(f"集合字段不可空，将重建集合：{collection_name}")
            utility.drop_collection(collection_name)
        # 检查并删除旧集合（如果结构不兼容）
        if utility.has_collection(collection_name):
            existing_collection = Collection(collection_name)
            return existing_collection

        # 定义集合结构，关键修改：设置nullable=True
        fields = [
            FieldSchema(name="id", dtype=DataType.VARCHAR, max_length=64, is_primary=True),
            FieldSchema(name=self.vector_field_name, dtype=DataType.FLOAT_VECTOR, dim=settings.MILVUS_DIM),
            FieldSchema(name="content", dtype=DataType.VARCHAR, max_length=5120),
            FieldSchema(name="file_id", dtype=DataType.VARCHAR, max_length=64),
            FieldSchema(name="file_name", dtype=DataType.VARCHAR, max_length=255),
            FieldSchema(name="chunk_id", dtype=DataType.VARCHAR, max_length=64),
            FieldSchema(name="file_ext", dtype=DataType.VARCHAR, max_length=100),
            FieldSchema(name="created_at", dtype=DataType.VARCHAR, max_length=32),
            FieldSchema(name="Comments",   dtype=DataType.VARCHAR,  max_length=2024, is_nullable=True),
            FieldSchema(name="Trapped", dtype=DataType.VARCHAR, max_length=2024, is_nullable=True),
        ]

        schema = CollectionSchema(
            fields=fields,
            description="合同知识库集合",
            enable_dynamic_field=False
        )

        collection = Collection(
            name=collection_name,
            schema=schema,
            using="default"
        )

        # 创建索引
        index_params = {
            "index_type": "IVF_FLAT",
            "metric_type": "L2",
            "params": {"nlist": 128}
        }
        collection.create_index(
            field_name=self.vector_field_name,
            index_params=index_params
        )

        logger.info(f"成功创建集合：{collection_name}（包含可空字段）")
        return collection

    def _init_langchain_milvus(self) -> MilvusVectorStore:
        """初始化LangChain Milvus向量存储，明确指定所有字段映射"""
        # 额外参数：确保LangChain不自动创建索引（我们已经手动创建）
        return MilvusVectorStore(
            collection_name=self.collection_name,
            embedding_function=self.embeddings,
            connection_args={
                "host": settings.MILVUS_HOST,
                "port": settings.MILVUS_PORT,
                "alias": "default"  # 明确指定连接别名
            },
            vector_field=self.vector_field_name,  # 向量字段
            text_field="content",  # 文本内容字段
            primary_field="id",  # 主键字段
            index_params={
                "index_type": "IVF_FLAT",
                "metric_type": "L2",
                "params": {"nlist": 128}
            },
            # 关键：禁用自动创建索引，使用我们手动创建的索引
            auto_id=False,
            drop_old=False
        )

    def add_documents_with_retry(self, docs: List[Document], ids: List[str], batch_size: int = 10) -> int:
        """带重试和分批处理的文档添加方法"""
        success_count = 0
        total = len(docs)

        # 补全字段
        # 补全缺失的字段，添加日志验证
        # 补全缺失的字段
        for doc in docs:
            doc.metadata.setdefault("Comments", "")
            doc.metadata.setdefault("Trapped", "")
            doc.metadata.setdefault("file_id", "")
            doc.metadata.setdefault("file_name", "")
            doc.metadata.setdefault("chunk_id", "")
            doc.metadata.setdefault("file_ext", "")
            doc.metadata.setdefault("created_at", time.strftime("%Y-%m-%dT%H:%M:%S"))
            # 记录当前内容长度
            doc.metadata["content_length"] = len(doc.page_content)


        for i in range(0, total, batch_size):
            batch_docs = docs[i:i + batch_size]
            batch_ids = ids[i:i + batch_size]
            batch_num = (i // batch_size) + 1

            for attempt in range(self.max_retries):
                try:
                    # 检查连接状态
                    if not connections.has_connection("default"):
                        logger.warning("Milvus连接已断开，尝试重连...")
                        self._connect()

                    self.vector_store.add_documents(
                        documents=batch_docs,
                        ids=batch_ids
                    )
                    success_count += len(batch_docs)
                    logger.info(f"第{batch_num}批入库成功（{len(batch_docs)}个分片）")
                    break
                except Exception as e:
                    logger.error(f"第{batch_num}批入库尝试 {attempt + 1}/{self.max_retries} 失败：{str(e)}")
                    if attempt < self.max_retries - 1:
                        time.sleep(self.retry_delay)
                    else:
                        logger.error(f"第{batch_num}批入库最终失败，将跳过")

        return success_count

    # 其他方法保持不变（字段名已在集合中统一为 'vector'，无需修改）
    def add_documents(self, docs: List[Document], file_id: str) -> Dict[str, Any]:
        try:
            doc_ids = [f"{file_id}_chunk_{i:04d}" for i in range(len(docs))]
            for i, doc in enumerate(docs):
                doc.metadata.update({
                    "file_id": file_id,
                    "chunk_id": doc_ids[i].split("_chunk_")[-1],
                    "created_at": doc.metadata.get("created_at", "")
                })
            self.vector_store.add_documents(documents=docs, ids=doc_ids)
            self.collection.flush()
            total_count = self.collection.num_entities
            logger.info(f"成功添加{len(docs)}个文档分片到Milvus，总数量：{total_count}")
            return {
                "status": "success",
                "added_count": len(docs),
                "total_count": total_count,
                "file_id": file_id
            }
        except Exception as e:
            logger.error(f"Milvus添加文档失败：{str(e)}")
            raise

    def search(self, query: str, file_id: Optional[str] = None, top_k: int = 5) -> List[Document]:
        filter_expr = f"file_id == '{file_id}'" if file_id else None
        results = self.vector_store.similarity_search(
            query=query, k=top_k, expr=filter_expr
        )
        logger.info(f"检索完成：查询='{query[:50]}...'，结果数={len(results)}")
        return results

    def delete_by_file_id(self, file_id: str) -> Dict[str, Any]:
        try:
            expr = f"file_id == '{file_id}'"
            delete_count = self.collection.delete(expr)["delete_count"]
            self.collection.flush()
            total_count = self.collection.num_entities
            logger.info(f"按file_id删除文档：{file_id}，删除数量：{delete_count}")
            return {
                "status": "success",
                "deleted_count": delete_count,
                "remaining_count": total_count,
                "file_id": file_id
            }
        except Exception as e:
            logger.error(f"Milvus删除文档失败：{str(e)}")
            raise


# 单例存储
# milvus_store = MilvusStoreWrapper()
